import pandas as pd
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset
from sklearn.preprocessing import StandardScaler
import matplotlib.pyplot as plt
from matplotlib import rcParams
import numpy as np
from matplotlib.ticker import ScalarFormatter

# 设置字体以支持中文
rcParams['font.sans-serif'] = ['SimHei']  # 使用黑体字体
rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

# 读取数据
data = pd.read_excel('../output.xlsx')

# 根据作物编号进行分类
grain_indices = data[(data['作物编号'] >= 1) & (data['作物编号'] <= 16)].index
vegetable_indices = data[(data['作物编号'] >= 17) & (data['作物编号'] <= 37)].index
fungi_indices = data[(data['作物编号'] >= 38) & (data['作物编号'] <= 41)].index

# 提取不同类别的数据
X_cost_grain = data.loc[grain_indices, ['种植成本']].values
Y_grain = data.loc[grain_indices, '预期销售量'].values

X_cost_vegetable = data.loc[vegetable_indices, ['种植成本']].values
Y_vegetable = data.loc[vegetable_indices, '预期销售量'].values

X_price_vegetable = data.loc[vegetable_indices, ['销售价格']].values

X_cost_fungi = data.loc[fungi_indices, ['种植成本']].values
X_price_fungi = data.loc[fungi_indices, ['销售价格']].values
Y_fungi = data.loc[fungi_indices, '预期销售量'].values

# 标准化数据
scaler_X_cost = StandardScaler()
scaler_X_price = StandardScaler()
scaler_Y = StandardScaler()

X_cost_grain = scaler_X_cost.fit_transform(X_cost_grain)
X_cost_vegetable = scaler_X_cost.fit_transform(X_cost_vegetable)
X_cost_fungi = scaler_X_cost.fit_transform(X_cost_fungi)

X_price_vegetable = scaler_X_price.fit_transform(X_price_vegetable)
X_price_fungi = scaler_X_price.fit_transform(X_price_fungi)

Y_grain = scaler_Y.fit_transform(Y_grain.reshape(-1, 1)).flatten()
Y_vegetable = scaler_Y.fit_transform(Y_vegetable.reshape(-1, 1)).flatten()
Y_fungi = scaler_Y.fit_transform(Y_fungi.reshape(-1, 1)).flatten()

# 转换为 PyTorch 张量
X_cost_grain_tensor = torch.tensor(X_cost_grain, dtype=torch.float32)
X_cost_vegetable_tensor = torch.tensor(X_cost_vegetable, dtype=torch.float32)
X_price_vegetable_tensor = torch.tensor(X_price_vegetable, dtype=torch.float32)
X_price_fungi_tensor = torch.tensor(X_price_fungi, dtype=torch.float32)
X_cost_fungi_tensor = torch.tensor(X_cost_fungi, dtype=torch.float32)

Y_grain_tensor = torch.tensor(Y_grain, dtype=torch.float32).view(-1, 1)
Y_vegetable_tensor = torch.tensor(Y_vegetable, dtype=torch.float32).view(-1, 1)
Y_fungi_tensor = torch.tensor(Y_fungi, dtype=torch.float32).view(-1, 1)

# 创建数据加载器
grain_cost_loader = DataLoader(TensorDataset(X_cost_grain_tensor, Y_grain_tensor), batch_size=16, shuffle=True)
vegetable_cost_loader = DataLoader(TensorDataset(X_cost_vegetable_tensor, Y_vegetable_tensor), batch_size=16, shuffle=True)
vegetable_price_loader = DataLoader(TensorDataset(X_price_vegetable_tensor, Y_vegetable_tensor), batch_size=16, shuffle=True)
fungi_price_loader = DataLoader(TensorDataset(X_price_fungi_tensor, Y_fungi_tensor), batch_size=16, shuffle=True)
fungi_cost_loader = DataLoader(TensorDataset(X_cost_fungi_tensor, Y_fungi_tensor), batch_size=16, shuffle=True)

# 神经网络模型
class SimpleMLP(nn.Module):
    def __init__(self):
        super(SimpleMLP, self).__init__()
        self.hidden1 = nn.Linear(1, 64)
        self.hidden2 = nn.Linear(64, 32)
        self.output = nn.Linear(32, 1)
        self.activation = nn.Softplus()

    def forward(self, x):
        x = torch.relu(self.hidden1(x))
        x = torch.relu(self.hidden2(x))
        x = self.activation(self.output(x))
        return x

# 训练和预测
def train_and_predict(loader, X_tensor):
    model = SimpleMLP()
    criterion = nn.MSELoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

    epochs = 500
    for epoch in range(epochs):
        model.train()
        for X_batch, Y_batch in loader:
            optimizer.zero_grad()
            Y_pred = model(X_batch)
            loss = criterion(Y_pred, Y_batch)
            loss.backward()
            optimizer.step()

    model.eval()
    with torch.no_grad():
        Y_pred = model(X_tensor).numpy()
    return model, Y_pred

# 训练模型并预测
grain_cost_model, Y_grain_cost_pred = train_and_predict(grain_cost_loader, X_cost_grain_tensor)
vegetable_cost_model, Y_vegetable_cost_pred = train_and_predict(vegetable_cost_loader, X_cost_vegetable_tensor)
vegetable_price_model, Y_vegetable_price_pred = train_and_predict(vegetable_price_loader, X_price_vegetable_tensor)
fungi_price_model, Y_fungi_price_pred = train_and_predict(fungi_price_loader, X_price_fungi_tensor)
fungi_cost_model, Y_fungi_cost_pred = train_and_predict(fungi_cost_loader, X_cost_fungi_tensor)

# 反标准化预测结果和实际值
Y_grain_cost_pred_actual = scaler_Y.inverse_transform(Y_grain_cost_pred.reshape(-1, 1)).flatten()
Y_vegetable_cost_pred_actual = scaler_Y.inverse_transform(Y_vegetable_cost_pred.reshape(-1, 1)).flatten()
Y_vegetable_price_pred_actual = scaler_Y.inverse_transform(Y_vegetable_price_pred.reshape(-1, 1)).flatten()
Y_fungi_price_pred_actual = scaler_Y.inverse_transform(Y_fungi_price_pred.reshape(-1, 1)).flatten()
Y_fungi_cost_pred_actual = scaler_Y.inverse_transform(Y_fungi_cost_pred.reshape(-1, 1)).flatten()

Y_grain_actual = scaler_Y.inverse_transform(Y_grain.reshape(-1, 1)).flatten()
Y_vegetable_actual = scaler_Y.inverse_transform(Y_vegetable.reshape(-1, 1)).flatten()
Y_fungi_actual = scaler_Y.inverse_transform(Y_fungi.reshape(-1, 1)).flatten()

# 生成并绘制拟合曲线
def generate_fit_curve(model, X_min, X_max, scaler_X, scaler_Y):
    X_smooth = np.linspace(X_min, X_max, 100).reshape(-1, 1)  # 生成平滑的100个输入点
    X_smooth_tensor = torch.tensor(scaler_X.transform(X_smooth), dtype=torch.float32)  # 标准化
    with torch.no_grad():
        Y_smooth_pred = model(X_smooth_tensor).numpy()  # 模型预测
    Y_smooth_pred_actual = scaler_Y.inverse_transform(Y_smooth_pred).flatten()  # 反标准化Y
    return X_smooth, Y_smooth_pred_actual

# 绘制图表
plt.figure(figsize=(12, 16))
formatter = ScalarFormatter(useOffset=False)

# 粮食种植成本 vs 预期销售量
plt.subplot(3, 2, 1)
plt.scatter(scaler_X_cost.inverse_transform(X_cost_grain), Y_grain_actual, color='blue', label='实际值')
plt.scatter(scaler_X_cost.inverse_transform(X_cost_grain), Y_grain_cost_pred_actual, color='red', label='预测值')
X_smooth_grain, Y_smooth_grain_actual = generate_fit_curve(grain_cost_model, scaler_X_cost.inverse_transform(X_cost_grain).min(), scaler_X_cost.inverse_transform(X_cost_grain).max(), scaler_X_cost, scaler_Y)
plt.plot(X_smooth_grain, Y_smooth_grain_actual, color='green', label='拟合曲线')
plt.xlabel('种植成本')
plt.ylabel('预期销售量')
plt.title('粮食种植成本 vs 预期销售量 拟合曲线')
plt.legend()

# 蔬菜种植成本 vs 预期销售量
plt.subplot(3, 2, 2)
plt.scatter(scaler_X_cost.inverse_transform(X_cost_vegetable), Y_vegetable_actual, color='blue', label='实际值')
plt.scatter(scaler_X_cost.inverse_transform(X_cost_vegetable), Y_vegetable_cost_pred_actual, color='red', label='预测值')
X_smooth_vegetable_cost, Y_smooth_vegetable_cost_actual = generate_fit_curve(vegetable_cost_model, scaler_X_cost.inverse_transform(X_cost_vegetable).min(), scaler_X_cost.inverse_transform(X_cost_vegetable).max(), scaler_X_cost, scaler_Y)
plt.plot(X_smooth_vegetable_cost, Y_smooth_vegetable_cost_actual, color='green', label='拟合曲线')
plt.xlabel('种植成本')
plt.ylabel('预期销售量')
plt.title('蔬菜种植成本 vs 预期销售量 拟合曲线')
plt.legend()

# 蔬菜销售价格 vs 预期销售量
plt.subplot(3, 2, 3)
plt.scatter(scaler_X_price.inverse_transform(X_price_vegetable), Y_vegetable_actual, color='blue', label='实际值')
plt.scatter(scaler_X_price.inverse_transform(X_price_vegetable), Y_vegetable_price_pred_actual, color='red', label='预测值')
X_smooth_vegetable_price, Y_smooth_vegetable_price_actual = generate_fit_curve(vegetable_price_model, scaler_X_price.inverse_transform(X_price_vegetable).min(), scaler_X_price.inverse_transform(X_price_vegetable).max(), scaler_X_price, scaler_Y)
plt.plot(X_smooth_vegetable_price, Y_smooth_vegetable_price_actual, color='green', label='拟合曲线')
plt.xlabel('销售价格')
plt.ylabel('预期销售量')
plt.title('蔬菜销售价格 vs 预期销售量 拟合曲线')
plt.legend()

# 食用菌销售价格 vs 预期销售量
plt.subplot(3, 2, 4)
plt.scatter(scaler_X_price.inverse_transform(X_price_fungi), Y_fungi_actual, color='blue', label='实际值')
plt.scatter(scaler_X_price.inverse_transform(X_price_fungi), Y_fungi_price_pred_actual, color='red', label='预测值')
X_smooth_fungi_price, Y_smooth_fungi_price_actual = generate_fit_curve(fungi_price_model, scaler_X_price.inverse_transform(X_price_fungi).min(), scaler_X_price.inverse_transform(X_price_fungi).max(), scaler_X_price, scaler_Y)
plt.plot(X_smooth_fungi_price, Y_smooth_fungi_price_actual, color='green', label='拟合曲线')
plt.xlabel('销售价格')
plt.ylabel('预期销售量')
plt.title('食用菌销售价格 vs 预期销售量 拟合曲线')
plt.legend()

# 食用菌种植成本 vs 预期销售量
plt.subplot(3, 2, 5)
plt.scatter(scaler_X_cost.inverse_transform(X_cost_fungi), Y_fungi_actual, color='blue', label='实际值')
plt.scatter(scaler_X_cost.inverse_transform(X_cost_fungi), Y_fungi_cost_pred_actual, color='red', label='预测值')
X_smooth_fungi_cost, Y_smooth_fungi_cost_actual = generate_fit_curve(fungi_cost_model, scaler_X_cost.inverse_transform(X_cost_fungi).min(), scaler_X_cost.inverse_transform(X_cost_fungi).max(), scaler_X_cost, scaler_Y)
plt.plot(X_smooth_fungi_cost, Y_smooth_fungi_cost_actual, color='green', label='拟合曲线')
plt.xlabel('种植成本')
plt.ylabel('预期销售量')
plt.title('食用菌种植成本 vs 预期销售量 拟合曲线')
plt.legend()

plt.subplots_adjust(left=0.1, right=0.95, top=0.95, bottom=0.2, hspace=1.0, wspace=0.35)

plt.show()
